"""
# arxiv的字段对应的内容
# entry_id A url http://arxiv.org/abs/{id}。
# updated 当结果最后一次更新的时间。
# published 结果最初发布的时间。
# title 结果的标题。
# authors 结果的作者，作为 arxiv.Authors。
# summary 结果的摘要。
# comment 如果有的话，作者的评论。
# journal_ref 如果有的话，期刊引用。
# doi 如果有的话，解析到外部资源的DOI的URL。
# primary_category 结果的主要arXiv类别。参见arXiv: 类别体系结构[4]。
# categories 结果的所有类别。参见arXiv: 类别体系结构。
# links 与此结果相关的最多三个URL，作为 arxiv.Links。
# pdf_url 如果有的话，结果的PDF的URL。注意：这个URL也出现在 result.links 中。
"""
from typing import Any
import asyncio
from LLM import llm
import arxiv
import time
from pydantic import BaseModel
from datetime import datetime

class ArxivResult(BaseModel):
    entry_id: str
    published: Any
    title: str
    authors: list[Any]
    summary: str
    pdf_url: str


class ArxivSearch:
    def __init__(self, **kwargs):
        self.llm = llm.LLM()
        self.logger = kwargs.get('logger')
    def arxiv_search(self, query, max_results):
        time_arxiv_start = time.time()
        client = arxiv.Client(page_size=50, delay_seconds=3.0, num_retries=3)
        search = arxiv.Search(
            query=query,
            max_results=max_results,
            sort_by=arxiv.SortCriterion.SubmittedDate,
        )
        result_list = []
        for result in client.results(search):
            result_obj = ArxivResult(entry_id=result.entry_id, published=result.published, title=result.title,
                                     authors=result.authors, summary=result.summary, pdf_url=result.pdf_url)
            result_list.append(result_obj)
        time_arxiv_end = time.time()
        print(time_arxiv_end - time_arxiv_start)
        return result_list

    async def run(self, article_info):
        error = True
        total_result = """这里是Agent Daily，每早为您汇总Agent论文。
        
        """
        for direction in article_info.directions:
            total_result += f"""
            <p style="font-size:24px;"><strong>{direction}方向：</strong></p>
"""
            result = self.arxiv_search(query=direction, max_results=article_info.max_results)
            prompt = "论文摘要如下：\n{arxiv_summary}\n\n请将摘要翻译成中文"
            tasks = []
            arxiv_result_list = []
            for res in result:
                if datetime.strptime(article_info.start_time.strftime('%Y-%m-%d'), '%Y-%m-%d') <= datetime.strptime(res.published.strftime('%Y-%m-%d'), '%Y-%m-%d') <= datetime.strptime(article_info.end_time.strftime('%Y-%m-%d'), '%Y-%m-%d'):
                    arxiv_result = f"""
                    <p><strong>论文标题：</strong><em>{res.title}</em></p>\r<p><strong>论文链接：</strong></p>\r<p><a href="{res.entry_id}" style="font-size:10px;">{res.entry_id}</a></p>\r<p><strong>论文发布日期：</strong>{str(res.published)[0:10]}</p>\r<p><strong>论文作者：</strong>{res.authors[0].name}</p>"""
                    arxiv_result_list.append(arxiv_result)
                    arxiv_summary = res.summary
                    prompt_result = prompt.format(arxiv_summary=arxiv_summary)
                    task = asyncio.create_task(self.llm.chatglm(prompt_result, [], 'glm-4-air'))
                    tasks.append(task)
            if tasks != []:
                error = False
                all_results = await asyncio.gather(*tasks)
                for i in range(len(all_results)):
                    simple_result = f"{arxiv_result_list[i]}" + f"<p><strong>论文概述：</strong>{all_results[i]}</p>"
                    total_result += simple_result
        title_message = """以上是今天的论文主要内容，你需要根据该内容分别在RAG、Agent和Multimodel中挑选一个重点论文进行15个字以内的一句话概括，直接输出45字内的中文结果。直接输出概括即可，不要输出Agent方向、RAG方向、Multimodel方向这种提示信息!
        示例：一种结合xxx的RAG方法，Agent的xxx在xxx规划方向的创新应用，多模态语言模型的xxx方法改进"""
        title = await self.llm.chatglm(total_result+title_message, [], 'glm-4')
        return total_result, title[0:45], error
